from typing import Dict
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.configuration import conf
from datetime import timedelta, timezone
import re
import MySQLdb


class DorisLoadOperator(BaseOperator):

    @apply_defaults
    def __init__(
            self,
            *,
            content: str,
            **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.content = content

    def execute(self, context: Dict):
        '''
        {{day}} {{next_day}}
        {{hour}} {{next_hour}}
        :param context:
        :return:
        '''

        if not self.content:
            print('content is none, exit')
            return

        execution_date = context['execution_date']
        next_execution_date = context['next_execution_date']
        day = execution_date.astimezone(timezone(timedelta(hours=8))).strftime("%Y%m%d")
        next_day = next_execution_date.astimezone(timezone(timedelta(hours=8))).strftime("%Y%m%d")
        hour = execution_date.astimezone(timezone(timedelta(hours=8))).strftime("%H")
        next_hour = next_execution_date.astimezone(timezone(timedelta(hours=8))).strftime("%H")

        self.content = re.sub(r'\{\{\s*day\s*\}\}', day, self.content)
        self.content = re.sub(r'\{\{\s*next_day\s*\}\}', next_day, self.content)
        self.content = re.sub(r'\{\{\s*hour\s*\}\}', hour, self.content)
        self.content = re.sub(r'\{\{\s*next_hour\s*\}\}', next_hour, self.content)

        host = conf.get('doris', 'host')
        port = int(conf.get('doris', 'port'))
        user = conf.get('doris', 'user')
        password = conf.get('doris', 'password')
        if not password:
            password = ''
        db = conf.get('doris', 'password')

        db = MySQLdb.connect(host, user, password, db, charset='utf8', port=port)
        cursor = db.cursor()
        cursor.connection.autocommit(True)

        cmds = self.content.split(';')
        for cmd in cmds:
            sql = cmd.strip()
            if sql:
                print('execute : {}'.format(sql))
                try:
                    cursor.execute(sql)
                except Exception as e:
                    print(e)
        db.close()
